# RBC教材第四章可变劳动的例子
# 该算法有几个预备步骤：
# 1. 构造贝尔曼方程

rm(list = ls())
devtools::load_all()
library(rootSolve)
library(signal)
library(magrittr)
library(ggplot2)
# 参数设置
deltax <- 0.1
thetax <- 0.36
betax <- 0.98
A <- 0.5

# 数值计算设置
crit <- 1
tol <- 0.001
eps <- 0.02

# 格点和值函数初值设置和下一期最优资本存量内存分配
k0 <- seq(0.06,10,length.out = 80)
v <- seq(-4,9,length.out = length(k0))
aopt <- rep(0,length(k0))

ans <- Sys.time()
j <- 0
while (crit > tol) {
  j <- j + 1
  print(crit)
  vold <- v
  m <- 0
  for (i in 1:length(k0)) {
    ainit <- k0[i]
    v0 <- -1e10
    m <- m + 1

    bellman <- para_cmp(a0 = ainit, agrid = k0, vold = vold,
             thetax = thetax, deltax = deltax, betax = betax, A = A)

    # 给定ainit，确定下一期资本存量的范围。注意此处a1取了m格点处的值
    # 意味着新计算的值函数,a1从m点开始
    pos <- decise_range(k0, bellman, m0 = m)
    # browser()
    m <- pos$m0 # 更新位置

    # 给定a1范围，找到考虑边界点的最大值函数对应的a1
    aopt[i] <- findmax(pos,eps,agrid = k0,vold, bellman)
    v[i] <- bellman(a1 = aopt[i])
  }
  crit <- abs(vold-v) %>% mean() %>% max()
}
Sys.time()- ans

# 获得劳动力
chfun <- function(h,a0,a1,A,thetax,deltax){
  c0 <- a0^thetax*h^(1-thetax) - a1 + (1-deltax)*a0
  A*c0/(1-h)-(1-thetax)*a0^thetax*h^(-thetax)
}

hh <- numeric(80)
for (i in 1:80) {
  hh[i] <- uniroot.all(chfun, interval = c(0,1), a0 = k0[i], a1 = aopt[i], A = A,
                       thetax = thetax, deltax = deltax)
}
# 获得消费
c0 <- k0^thetax*hh^(1-thetax) - aopt + (1-deltax)*k0


picdata <- cbind(k0,aopt,hh,c0,v) %>% as.data.frame()
ggplot(picdata, aes(x = k0, y = aopt)) + geom_line() +
  geom_line(aes(y = k0), color = 'red')

